{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    },
    "colab": {
      "name": "_graph_optimization.ipynb",
      "provenance": [],
      "private_outputs": true,
      "toc_visible": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYLyuStTYesc",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PVm-iEoxYesf",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3MPf91rVYesq",
        "colab_type": "text"
      },
      "source": [
        "# TensorFlow Graph Optimization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zmNCsZlgYesr",
        "colab_type": "text"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/not_a_real_link\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/tools/templates/notebook.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/tools/templates/notebook.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/tools/templates/notebook.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l0qacLgyYess",
        "colab_type": "text"
      },
      "source": [
        "## Overview\n",
        "\n",
        "TensorFlow optimizes the TensorFlow graph (built by `tf.function`) before executing operations.\n",
        "This makes the TensorFlow graph more efficiently and less peak memory usage.\n",
        "As most of optimizations are enabled automatically, you don't need to know what optimizations are perfomred under the TensorFlow.\n",
        "\n",
        "On the other hand, TensorFlow provides the way for the advanced TensorFlow users to enable/disable the optimization by using `tf.config.optimizer.set_experimental_options()`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZAUsxyWYess",
        "colab_type": "text"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sQ0HvpqDYest",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "\n",
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x  #gpu\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-zkJgR5Yesw",
        "colab_type": "text"
      },
      "source": [
        "## Graph Optimization performed in TensorFlow\n",
        "\n",
        "TensorFlow performs many optimizations to the TensorFlow Graph such as Constant Foling, Memory Optimizer, etc...  \n",
        "All optimizations are listed in [tf.config.optimizer.set_experimental_options()](https://www.tensorflow.org/api_docs/python/tf/config/optimizer/set_experimental_options) API document page.\n",
        "Below optimizations are typical optimizations performed in TensorFlow.\n",
        "\n",
        "* **Constant Folding** evaluates the operations whose input tensors are all constant.\n",
        "* **Memory Optimizer** analyzes the TensorFlow graph to inspect the peak memory usage at each operations. To reduce a peak memory usage, Memory Optimizer inserts CPU-GPU memory copy operations for swapping GPU memory to CPU, or rewrites TensorFlow graph for the recomputation.\n",
        "* **Layout Optimizer** inserts Transpose operation to change data-format in order to execute data-format depended operation more efficiently on GPU.\n",
        "* **Arithmentic Optimizer** is an optimization to rewrite the TensorFlow graph using mathematical equivalence relation.\n",
        "* **Debug Stripper** strips debug operations such as Assert.\n",
        "\n",
        "More detail information about the graph optimization can be found in [this slide](http://web.stanford.edu/class/cs245/slides/TFGraphOptimizationsStanford.pdf)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kx92OO9cYesy",
        "colab_type": "text"
      },
      "source": [
        "## Checking Enabled Graph Optimization\n",
        "\n",
        "Enabled optimization configuration can be shown by calling `tf.config.optimizer.get_experimental_options()`.\n",
        "\n",
        "Below simple code shows the default optimization configuration."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uRuhVoAlYesz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "tf.config.optimizer.get_experimental_options()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qQQj5YA8Yes1",
        "colab_type": "text"
      },
      "source": [
        "## Inspect the Optimized Graph\n",
        "\n",
        "The optimized graphs can be inspectable with using TensorBoard.\n",
        "\n",
        "Let us look at the case of \"Constant Folding\" optimization.\n",
        "\"Constant Folding\" is enabled by default, so you will see the effects of the constant propagations without any procedures.\n",
        "\n",
        "To output TensorBoard summary data, you must call `tf.summary.trace_on()` before executing the graph.\n",
        "All profile and graph data (includes optimized graphs) are stored in memory.\n",
        "After the graph execution, TensorBoard summary data can be output by calling `tf.summary.trace_export()`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GNy7cbz9Yes2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "@tf.function\n",
        "def simple_func(arg):\n",
        "    a = tf.constant(7.9)\n",
        "    b = tf.constant(6.3)\n",
        "    c = arg + a\n",
        "    d = a * b\n",
        "    ret = c + d\n",
        "    \n",
        "    return ret\n",
        "\n",
        "# Enable tracing data to inspect the optimized graph.\n",
        "writer = tf.summary.create_file_writer(\"summary\")\n",
        "tf.summary.trace_on(graph=True, profiler=True)\n",
        "\n",
        "arg = tf.constant(8.9)\n",
        "print(simple_func(arg))\n",
        "\n",
        "# Output traced data as TensorBoard summary data.\n",
        "with writer.as_default():\n",
        "    tf.summary.trace_export(\"summary\", step=0, profiler_outdir=\"summary\")\n",
        "\n",
        "# Disable tracing data.\n",
        "tf.summary.trace_off()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WD1-APylYes4",
        "colab_type": "text"
      },
      "source": [
        "Then, launch TensorBoard to inspect the optimized graph.\n",
        "\n",
        "The optimized graph is inspectable via Profile graph.\n",
        "The first image shows the user defined graph, and the second image shows the optimized graph.\n",
        "We can notice that mul node and it's input (Const) nodes are missing in the optimized graph.\n",
        "This indicates that Constant Folding works well against the user defined graph.\n",
        "\n",
        "![User Defined Graph](https://github.com/nuka137/docs/blob/graph_optimization/site/en/guide/images/graph_optimization/user_defined_graph.png?raw=1)\n",
        "\n",
        "![Optimized Graph](https://github.com/nuka137/docs/blob/graph_optimization/site/en/guide/images/graph_optimization/optimized_graph.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RmzHjejOYes5",
        "colab_type": "text"
      },
      "source": [
        "## Enable/Disable Graph Optimization\n",
        "\n",
        "Most of graph optimizations can be enabled/disabled by calling `tf.config.optimizer.set_experimental_options()`.\n",
        "\n",
        "Below simple code shows how to enable/disable the optimization."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-2EaOqmYes6",
        "colab_type": "text"
      },
      "source": [
        "### Disable \"Debug Stripper\"\n",
        "\n",
        "**Debug Stripper** which is disabled by default, strips the operations used for the debug purpose (Assert, CheckNumerics, Print).\n",
        "Below code will raise an exception at `tf.Assert`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NVpD71S1Yes8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "@tf.function\n",
        "def assert_func():\n",
        "    a = tf.constant(1.2)\n",
        "    computation_graph = tf.Assert(tf.less_equal(a, 1.0), [a])   # Will raise an \"InvalidArgumentError\" exception.\n",
        "    return a\n",
        "\n",
        "print(assert_func())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jTX6Z1MNYes-",
        "colab_type": "text"
      },
      "source": [
        "If you make Debug Stripper enabled, below code does not raise any exceptions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fxQjZu82Yes_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Enable \"Debug Stripper\".\n",
        "tf.config.optimizer.set_experimental_options({'debug_stripper': True})\n",
        "\n",
        "@tf.function\n",
        "def assert_func():\n",
        "    a = tf.constant(1.2)\n",
        "    computation_graph = tf.Assert(tf.less_equal(a, 1.0), [a])   # No exceptions are raised.\n",
        "    return a\n",
        "\n",
        "print(assert_func())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BY_L9lu-YetB",
        "colab_type": "text"
      },
      "source": [
        "### Disable All Graph Optimizations\n",
        "\n",
        "All Graph Optimization can be disabled when `'disable_meta_optimizer': False` is passed to `tf.config.optimizer.set_experimental_options()`.\n",
        "\n",
        "At first, let us check the optimized graph with default optimization configuration.\n",
        "Below code builds the graph which has a seqence of Transpose operations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SePP0NEkYetC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "@tf.function\n",
        "def optimized(arg):\n",
        "    a = arg * 2\n",
        "\n",
        "    # Will be simplified by Arithmetic Optimizer.\n",
        "    b = tf.transpose(a, perm=[1, 0])\n",
        "    ret = tf.transpose(b, perm=[1, 0])\n",
        "\n",
        "    return ret\n",
        "\n",
        "writer = tf.summary.create_file_writer(\"summary\")\n",
        "tf.summary.trace_on(graph=True, profiler=True)\n",
        "\n",
        "arg = tf.constant(np.random.normal(size=(30, 40)))\n",
        "optimized(arg)\n",
        "\n",
        "with writer.as_default():\n",
        "    tf.summary.trace_export(\"summary\", step=0, profiler_outdir=\"summary\")\n",
        "\n",
        "tf.summary.trace_off()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4TGgrQ0vYetF",
        "colab_type": "text"
      },
      "source": [
        "As you can see in TensorBoard, Transpose operation was erased by Arithmetic Optimizer to simplify Transpose operation into NoOp.\n",
        "\n",
        "![Meta Optimizer Enabled](https://github.com/nuka137/docs/blob/graph_optimization/site/en/guide/images/graph_optimization/meta_optimizer_enabled.png?raw=1)\n",
        "\n",
        "Next, let us check the graph with `'disable_meta_optimizer'` on."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P3pFWhdyYetG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "# Disable all graph optimizations.\n",
        "tf.config.optimizer.set_experimental_options({'disable_meta_optimizer': True})\n",
        "\n",
        "@tf.function\n",
        "def not_optimized(arg):\n",
        "    a = arg * 2\n",
        "    b = tf.transpose(a, perm=[1, 0])\n",
        "    ret = tf.transpose(b, perm=[1, 0])\n",
        "\n",
        "    return ret\n",
        "\n",
        "writer = tf.summary.create_file_writer(\"summary\")\n",
        "tf.summary.trace_on(graph=True, profiler=True)\n",
        "\n",
        "arg = tf.constant(np.random.normal(size=(30, 40)))\n",
        "not_optimized(arg)\n",
        "\n",
        "with writer.as_default():\n",
        "    tf.summary.trace_export(\"summary\", step=0, profiler_outdir=\"summary\")\n",
        "\n",
        "tf.summary.trace_off()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QKO5us9gYetJ",
        "colab_type": "text"
      },
      "source": [
        "In this graph, Transpose operations are remained in the graph, which indicates Arithmetic Optimizer was disabled.\n",
        "\n",
        "![Meta Optimizer Disabled](https://github.com/nuka137/docs/blob/graph_optimization/site/en/guide/images/graph_optimization/meta_optimizer_disabled.png?raw=1)"
      ]
    }
  ]
}